import os
from flask import Flask, render_template, request, jsonify, send_file
import cv2
import numpy as np

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB限制

# 确保上传目录存在
if not os.path.exists(app.config['UPLOAD_FOLDER']):
    os.makedirs(app.config['UPLOAD_FOLDER'])

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/lsb')
def lsb_page():
    return render_template('lsb.html')

@app.route('/dct')
def dct_page():
    return render_template('dct.html')

@app.route('/pvd')
def pvd_page():
    return render_template('pvd.html')

# 为每个算法创建独立的路由端点
@app.route('/embed_lsb', methods=['POST'])
def embed_lsb():
    return embed_message('lsb')

@app.route('/embed_dct', methods=['POST'])
def embed_dct():
    return embed_message('dct')

@app.route('/embed_pvd', methods=['POST'])
def embed_pvd():
    return embed_message('pvd')

@app.route('/extract_lsb', methods=['POST'])
def extract_lsb():
    return extract_message('lsb')

@app.route('/extract_dct', methods=['POST'])
def extract_dct():
    return extract_message('dct')

@app.route('/extract_pvd', methods=['POST'])
def extract_pvd():
    return extract_message('pvd')

def embed_message(algorithm):
    try:
        if 'image' not in request.files:
            return jsonify({'error': '没有选择图片'}), 400
        
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': '没有选择图片'}), 400
        
        message = request.form.get('message', '')
        if not message:
            return jsonify({'error': '没有输入隐藏信息'}), 400
        
        # 生成唯一文件名，强制使用PNG格式
        import uuid
        unique_filename = str(uuid.uuid4()) + '.png'
        filename = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
        file.save(filename)
        
        print(f"处理文件: {filename}, 消息: {message}, 算法: {algorithm}")
        
        # 根据算法选择处理方式
        if algorithm == 'lsb':
            from algorithms.lsb import LSB
            processor = LSB()
        elif algorithm == 'dct':
            from algorithms.dct import DCT
            processor = DCT()
        elif algorithm == 'pvd':
            from algorithms.pvd import PVD
            processor = PVD()
        else:
            return jsonify({'error': '不支持的算法'}), 400
        
        # 嵌入信息
        output_filename = 'embedded_' + unique_filename
        output_path = os.path.join(app.config['UPLOAD_FOLDER'], output_filename)
        success = processor.embed(filename, message, output_path)
        
        # 清理原始上传文件
        if os.path.exists(filename):
            os.remove(filename)
        
        if success:
            return jsonify({
                'success': True,
                'message': '信息隐藏成功',
                'image_url': f'/download/{output_filename}'
            })
        else:
            # 清理失败的文件
            if os.path.exists(output_path):
                os.remove(output_path)
            return jsonify({'error': '信息隐藏失败，请检查图像格式和消息长度'}), 500
            
    except Exception as e:
        print(f"嵌入路由错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

def extract_message(algorithm):
    try:
        if 'image' not in request.files:
            return jsonify({'error': '没有选择图片'}), 400
        
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': '没有选择图片'}), 400
        
        # 保存上传的图片
        import uuid
        unique_filename = str(uuid.uuid4()) + '.png'
        filename = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
        file.save(filename)
        
        print(f"提取文件: {filename}, 算法: {algorithm}")
        
        # 根据算法选择处理方式
        if algorithm == 'lsb':
            from algorithms.lsb import LSB
            processor = LSB()
        elif algorithm == 'dct':
            from algorithms.dct import DCT
            processor = DCT()
        elif algorithm == 'pvd':
            from algorithms.pvd import PVD
            processor = PVD()
        else:
            return jsonify({'error': '不支持的算法'}), 400
        
        # 提取信息
        extracted_message = processor.extract(filename)
        
        # 清理上传文件
        if os.path.exists(filename):
            os.remove(filename)
        
        return jsonify({
            'success': True,
            'message': extracted_message
        })
            
    except Exception as e:
        print(f"提取路由错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/download/<filename>')
def download_file(filename):
    return send_file(os.path.join(app.config['UPLOAD_FOLDER'], filename), as_attachment=True)

if __name__ == '__main__':
    app.run(debug=True)